import torch
from torch import optim
from torch.utils.data import DataLoader
from torchvision import transforms

from dateSet import Dataset
from deepSupervision import MultipleOutputLoss2
from loss import UNetPlusPlusLoss
from model import NestedUNet


def train(model, dataloader, criterion, optimizer, device, num_epochs=31):
    model.train()
    for epoch in range(num_epochs):
        print(f'第{epoch + 1}轮训练:')
        for images, masks, _ in dataloader:
            images = images.float().to(device)
            masks = masks.float().to(device)

            optimizer.zero_grad()
            outputs = model(images)
            # 对每一个output计算loss
            loss_list = []
            for output in outputs:
                loss_list.append(criterion.cul_loss(output, masks))
            loss = MultipleOutputLoss2([1, 1, 1, 1]).cul_deep_loss(loss_list)
            loss.backward()
            optimizer.step()
            outputs.clear()

        print(f'Epoch {epoch}, Loss: {loss.item()}')
        if (epoch + 1) % 3 == 0:
            torch.save(model.state_dict(), f'UPP_Cell-{epoch}.pt')


if __name__ == '__main__':
    dataset = Dataset(
        root_dir='/Volumes/For_Mac/unet++/unet++/inputs/stage1_train',
        target_size=256,
        min_size=64,
        max_size=1024
    )

    # 数据集与数据加载器
    # 可以根据需求进行数据增强操作
    transform = transforms.Compose([
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])

    dataloader = DataLoader(dataset, batch_size=10, shuffle=True)

    # 初始化模型、损失函数和优化器
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = NestedUNet(1, 4, True).to(device)
    criterion = UNetPlusPlusLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)

    # 开始训练
    train(model, dataloader, criterion, optimizer, device)
